"""
# Preprocessing data

Preprocessing the images involves two steps:

- Resizing the image: Images are resized such that the shortest size is equal
to 800 px, after resizing if the longest side of the image exceeds 1333 px,
the image is resized such that the longest size is now capped at 1333 px.
- Applying augmentation: Random scale jittering  and random horizontal flipping
are the only augmentations applied to the images.

Along with the images, bounding boxes are rescaled and flipped if required.
"""

import tensorflow as tf


def random_flip_horizontal(image, boxes):
    """Flips image and boxes horizontally with 50% chance

    Args:
        image: A 3-D tensor of shape `(height, width, channels)` representing an
            image.
        boxes: A tensor with shape `(num_boxes, 4)` representing bounding boxes,
            having normalized coordinates.

    Returns:
      Randomly flipped image and boxes
    """
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_left_right(image)
        boxes = tf.stack(
            [1 - boxes[:, 2], boxes[:, 1], 1 - boxes[:, 0], boxes[:, 3]], axis=-1
        )
    return image, boxes


def resize_and_pad_image(
        image, min_side=256.0, max_side=512.0, jitter=[240, 384], stride=32.0):
    """Resizes and pads image while preserving aspect ratio.

    1. Resizes images so that the shorter side is equal to `min_side`
    2. If the longer side is greater than `max_side`, then resize the image
        with longer side equal to `max_side`
    3. Pad with zeros on right and bottom to make the image shape divisible by
        `stride`

    Args:
        image: A 3-D tensor of shape `(height, width, channels)` representing an
            image.
        min_side: The shorter side of the image is resized to this value, if
            `jitter` is set to None.
        max_side: If the longer side of the image exceeds this value after
            resizing, the image is resized such that the longer side now equals to
            this value.
        jitter: A list of floats containing minimum and maximum size for scale
            jittering. If available, the shorter side of the image will be
            resized to a random value in this range.
        stride: The stride of the smallest feature map in the feature pyramid.
            Can be calculated using `image_size / feature_map_size`.

    Returns:
        image: Resized and padded image.
        image_shape: Shape of the image before padding.
        ratio: The scaling factor used to resize the image
    """
    image_shape = tf.cast(tf.shape(image)[:2], dtype=tf.float32)
    if jitter is not None:
        min_side = tf.random.uniform(
            (), jitter[0], jitter[1], dtype=tf.float32)
    ratio = min_side / tf.reduce_min(image_shape)
    if ratio * tf.reduce_max(image_shape) > max_side:
        ratio = max_side / tf.reduce_max(image_shape)
    image_shape = ratio * image_shape
    padded_image_shape = tf.cast(tf.math.ceil(image_shape / stride) * stride,
                                 dtype=tf.int32)
    image = tf.image.resize(image, tf.cast(image_shape, dtype=tf.int32))
    image = tf.image.pad_to_bounding_box(image, 0, 0, padded_image_shape[0],
                                         padded_image_shape[1])

    return image, image_shape, ratio
